// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/core/simple_function_metadata.h>
#include <pollux/expression/signature_binder.h>
#include <pollux/expression/simple_function_adapter.h>
#include <pollux/type/type.h>

namespace kumo::pollux::exec {
    template<typename T>
    const std::shared_ptr<const T> &singleton_udf_metadata(
        bool defaultNullBehavior,
        const std::vector<exec::SignatureVariable> &constraints = {}) {
        static auto instance =
                std::make_shared<const T>(defaultNullBehavior, constraints);
        return instance;
    }

    using Function = SimpleFunctionAdapterFactory;
    using Metadata = core::ISimpleFunctionMetadata;
    using FunctionFactory = std::function<std::unique_ptr<Function>()>;

    struct FunctionEntry {
        FunctionEntry(
            const std::shared_ptr<const Metadata> &metadata,
            const FunctionFactory &factory)
            : metadata_{metadata}, factory_{factory} {
        }

        const Metadata &getMetadata() const {
            return *metadata_;
        }

        std::unique_ptr<Function> createFunction() const {
            return factory_();
        }

    private:
        const std::shared_ptr<const Metadata> metadata_;
        const FunctionFactory factory_;
    };

    using SignatureMap = std::unordered_map<
        FunctionSignature,
        std::vector<std::unique_ptr<const FunctionEntry> > >;
    using FunctionMap = std::unordered_map<std::string, SignatureMap>;

    class SimpleFunctionRegistry {
    public:
        /// Register a UDF with the given aliases and constraints. If an alias already
        /// exists and 'overwrite' is true, the existing entry in the function
        /// registry is overwritten by the current UDF. If an alias already exists and
        /// 'overwrite' is false, the current UDF is not registered with this alias.
        /// This method returns true if all 'aliases' are registered successfully. It
        /// returns false if any alias in 'aliases' already exists in the registry and
        /// is not overwritten.
        template<typename UDF>
        bool register_function(const std::vector<std::string> &aliases,
                              const std::vector<exec::SignatureVariable> &constraints, bool overwrite) {
            const auto &metadata = singleton_udf_metadata<typename UDF::Metadata>(UDF::is_default_null_behavior, constraints);
            const auto factory = []() { return CreateUdf<UDF>(); };

            if (aliases.empty()) {
                return register_function_internal(
                    metadata->getName(), metadata, factory, overwrite);
            } else {
                bool registered = true;
                for (const auto &name: aliases) {
                    registered &=
                            register_function_internal(name, metadata, factory, overwrite);
                }
                return registered;
            }
        }

        void removeFunction(const std::string &name);

        std::vector<std::string> get_function_names() const {
            std::vector<std::string> result;
            registeredFunctions_.withRLock([&](const auto &map) {
                result.reserve(map.size());

                for (const auto &entry: map) {
                    result.push_back(entry.first);
                }
            });
            return result;
        }

        void clear_registry() {
            registeredFunctions_.withWLock([&](auto &map) { map.clear(); });
        }

        std::vector<const FunctionSignature *> get_function_signatures(
            const std::string &name) const;

        std::vector<std::pair<VectorFunctionMetadata, const FunctionSignature *> >
        get_function_signatures_and_metadata(const std::string &name) const;

        class ResolvedSimpleFunction {
        public:
            ResolvedSimpleFunction(
                const FunctionEntry &functionEntry,
                const TypePtr &type)
                : functionEntry_(functionEntry), type_(type) {
            }

            auto createFunction() {
                return functionEntry_.createFunction();
            }

            const TypePtr &type() const {
                return type_;
            }

            std::string helpMessage(const std::string &name) const {
                return functionEntry_.getMetadata().helpMessage(name);
            }

            std::string toDebugString() const {
                return functionEntry_.getMetadata().toDebugString();
            }

            VectorFunctionMetadata metadata() const {
                return VectorFunctionMetadata{
                    false,
                    functionEntry_.getMetadata().isDeterministic(),
                    functionEntry_.getMetadata().defaultNullBehavior()
                };
            }

        private:
            const FunctionEntry &functionEntry_;
            const TypePtr type_;
        };

        std::optional<ResolvedSimpleFunction> resolve_function(
            const std::string &name,
            const std::vector<TypePtr> &argTypes) const;

    private:
        template<typename T>
        static std::unique_ptr<T> CreateUdf() {
            return std::make_unique<T>();
        }

        /// Registers a function with the given name and metadata. If an entry with
        /// the name already exists and 'overwrite' is true, the existing entry is
        /// overwritten. If an entry with the name already exists and 'overwrite' is
        /// false, the function reigstry remain unchanged. This method returns true if
        /// the function is successfully registered. It returns false if the function
        /// registry remains unchanged.
        bool register_function_internal(
            const std::string &name,
            const std::shared_ptr<const Metadata> &metadata,
            const FunctionFactory &factory,
            bool overwrite);

        melon::Synchronized<FunctionMap> registeredFunctions_;
    };

    const SimpleFunctionRegistry &simple_functions();

    SimpleFunctionRegistry &mutable_simple_functions();

    /// This function should be called once and alone.
    /// @param names Aliases for the function.
    /// @param constraints Additional constraints for variables used in function
    /// signature. Primarily used to specify rules for calculating precision and
    /// scale for decimal result types.
    /// @param overwrite If true, overwrites existing entries in the function
    /// registry with the same names.
    template<typename UDFHolder>
    bool register_simple_function(const std::vector<std::string> &names,
                                const std::vector<exec::SignatureVariable> &constraints, bool overwrite) {
        return mutable_simple_functions().register_function<SimpleFunctionAdapterFactoryImpl<UDFHolder> >(
            names, constraints, overwrite);
    }
} // namespace kumo::pollux::exec
